"""
数据库迁移脚本 - 添加用户权限管理功能
"""
import sqlite3
import os
from datetime import datetime
from passlib.context import CryptContext

# 密码加密
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

def get_password_hash(password: str) -> str:
    """获取密码哈希"""
    return pwd_context.hash(password)

def migrate_user_management():
    """迁移数据库，添加用户权限管理功能"""
    db_path = "users.db"
    
    if not os.path.exists(db_path):
        print("数据库文件不存在，跳过迁移")
        return
    
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    try:
        # 检查是否已经存在role字段
        cursor.execute("PRAGMA table_info(users)")
        columns = [column[1] for column in cursor.fetchall()]
        
        # 添加role字段到users表
        if 'role' not in columns:
            cursor.execute("ALTER TABLE users ADD COLUMN role TEXT DEFAULT 'user'")
            print("添加 role 字段到 users 表")
        
        # 创建login_records表
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS login_records (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                login_time DATETIME DEFAULT CURRENT_TIMESTAMP,
                ip_address TEXT,
                user_agent TEXT,
                status TEXT DEFAULT 'success',
                FOREIGN KEY (user_id) REFERENCES users (id)
            )
        """)
        print("创建 login_records 表")
        
        # 创建data_processing_records表
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS data_processing_records (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                operation_type TEXT NOT NULL,
                operation_details TEXT,
                file_name TEXT,
                processing_time DATETIME DEFAULT CURRENT_TIMESTAMP,
                duration_ms INTEGER,
                status TEXT DEFAULT 'success',
                error_message TEXT,
                FOREIGN KEY (user_id) REFERENCES users (id)
            )
        """)
        print("创建 data_processing_records 表")
        
        # 创建索引
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_login_records_user_id ON login_records(user_id)")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_login_records_login_time ON login_records(login_time)")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_processing_records_user_id ON data_processing_records(user_id)")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_processing_records_processing_time ON data_processing_records(processing_time)")
        print("创建索引")
        
        # 检查是否已存在admin用户
        cursor.execute("SELECT * FROM users WHERE username = 'admin'")
        admin_user = cursor.fetchone()
        
        if not admin_user:
            # 创建默认管理员账户
            admin_password_hash = get_password_hash("admin")
            cursor.execute("""
                INSERT INTO users (username, email, phone, hospital, position, hashed_password, role, is_active, created_at, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                "admin",
                "admin@unionstrong.com",
                "000-0000-0000",
                "系统管理",
                "主任",
                admin_password_hash,
                "admin",
                True,
                datetime.utcnow(),
                datetime.utcnow()
            ))
            print("创建默认管理员账户: admin / admin")
        else:
            # 如果admin用户已存在，更新其信息以符合验证要求
            cursor.execute("""
                UPDATE users SET 
                    role = 'admin',
                    email = 'admin@unionstrong.com',
                    position = '主任'
                WHERE username = 'admin'
            """)
            print("更新现有admin用户信息以符合验证要求")
        
        conn.commit()
        print("用户权限管理数据库迁移完成")
        
    except Exception as e:
        print(f"数据库迁移失败: {e}")
        conn.rollback()
    finally:
        conn.close()

if __name__ == "__main__":
    migrate_user_management() 